import numpy as np
import gzip
import os

# 检查并创建目录
if not os.path.exists('./dataset/mnist'):
    os.makedirs('./dataset/mnist')

def load_mnist_images(image_file):
    with gzip.open(image_file, 'rb') as f:
        data = np.frombuffer(f.read(), np.uint8, offset=16)
    return data.reshape(-1, 28, 28)

def load_mnist_labels(label_file):
    with gzip.open(label_file, 'rb') as f:
        data = np.frombuffer(f.read(), np.uint8, offset=8)
    return data

# 加载训练数据
x_train = load_mnist_images('./MNIST/raw/train-images-idx3-ubyte.gz')
y_train = load_mnist_labels('./MNIST/raw/train-labels-idx1-ubyte.gz')

# 加载测试数据
x_test = load_mnist_images('./MNIST/raw/t10k-images-idx3-ubyte.gz')
y_test = load_mnist_labels('./MNIST/raw/t10k-labels-idx1-ubyte.gz')

# 保存为 .npy 文件
np.save('./dataset/mnist/x_train.npy', x_train)
# 修改保存文件名
np.save('./dataset/mnist/y_train_label.npy', y_train)
np.save('./dataset/mnist/x_test.npy', x_test)
np.save('./dataset/mnist/y_test.npy', y_test)
